import os
import re
from decimal import Decimal, InvalidOperation, getcontext

from langchain_chroma import Chroma
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_openai import ChatOpenAI
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

#初始化通义千问模型-qwen-max
llm = ChatOpenAI(
    api_key="sk-b5a1d95d5595453dbafb47db0a80d4be",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model="qwen-max"
)

#初始化Embedding模型
embedding = DashScopeEmbeddings(
    model="text-embedding-v2",  # 阿里云提供的embedding模型
    dashscope_api_key=os.getenv("DASHSCOPE_API_KEY", "sk-b5a1d95d5595453dbafb47db0a80d4be")  # 使用环境变量或默认API key
)

#初始化向量数据库
db = Chroma(embedding_function=embedding,
            collection_metadata={"hnsw:space": "cosine"},
            persist_directory="./chroma_db")

def load_document_and_split(file_path):
    """加载文档并分割成文本块"""
    loader = TextLoader(file_path, encoding='utf-8')
    documents = loader.load()
    print(f"成功加载文档，共 {len(documents)} 个文档")
    
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len,
        separators=["\n\n", "\n", "。", "！", "？", ";", ":", "，", " ", ""],
        keep_separator=True,
        is_separator_regex=False,
        strip_whitespace=True,
    )
    
    splits = text_splitter.split_documents(documents)
    print(f"文档分割完成，共 {len(splits)} 个文本块")
    return documents, splits

def extract_rider_cost_info(content):
    """从内容中提取骑手成本信息，返回 {rider_id: {...}}"""
    cost_info = {}
    getcontext().prec = 50

    # 更健壮的正则：支持全角/半角括号，可选逗号与空格
    # 示例：1.骑手50016（S:50016Huarongwang），总成本166.68213851308604
    pattern = r"(\d+)\.骑手(\d+)[（(]([^）)]+)[）)],?\s*总成本\s*([0-9]+(?:\.[0-9]+)?(?:[eE][+-]?\d+)?)"
    
    print(f"开始解析成本信息，内容长度: {len(content)}")
    
    matches = list(re.finditer(pattern, content))
    print(f"找到 {len(matches)} 个匹配项")
    
    if not matches:
        # 调试：显示一些样本内容
        sample = content[:2000] if len(content) > 2000 else content
        print(f"未找到匹配，显示前2000字符样本内容：\n{sample}")
    
    for match in matches:
        rank, rider_id, rider_code, cost_str = match.groups()
        print(f"匹配到骑手: rank={rank}, rider_id={rider_id}, rider_code={rider_code}, cost_str={cost_str}")
        try:
            cost = Decimal(cost_str)
            cost_info[rider_id] = {
                'rank': int(rank),
                'rider_code': rider_code,
                'cost': cost,
                'cost_str': cost_str
            }
        except InvalidOperation:
            # 跳过无法解析的成本
            print(f"无法解析成本 {cost_str}")
            continue

    print(f"最终解析到 {len(cost_info)} 个骑手的成本信息")
    return cost_info

def analyze_rider_cost(user_question, file_path="e:\\llm\\03-langchain-rag\\data\\CA7044794619转换后.txt"):
    """
    分析骑手成本相关的问题 - 先收集完整信息再总结
    """
    try:
        # 加载文档（用于其他功能日志）
        documents, _ = load_document_and_split(file_path)
        
        # 从用户问题中提取骑手ID
        rider_id_match = re.search(r'\d+', user_question)
        if not rider_id_match:
            return "未能从问题中提取到骑手ID"
        rider_id = rider_id_match.group()
        print(f"检测到目标骑手ID: {rider_id}")

        # 直接读取原始文件文本，避免分块问题
        try:
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                raw_text = f.read()
        except Exception as e:
            return f"读取原始文件失败: {e}"

        # 1. 收集目标骑手的成本信息
        cost_info = extract_rider_cost_info(raw_text)
        target_cost_info = None
        if cost_info and rider_id in cost_info:
            target_cost_info = cost_info[rider_id]
            print(f"找到目标骑手成本信息: {target_cost_info}")

        # 2. 收集目标骑手的排除信息
        exclusion_pattern = f"订单ID35863807排除骑手{rider_id}\\(([^)]+)\\)，过滤器：([^，]+)，描述：([^，]+)，结果：([^，]+)，原因：([^。]+)"
        exclusion_match = re.search(exclusion_pattern, raw_text)
        exclusion_info = None
        if exclusion_match:
            rider_code, filter_type, description, result, reason = exclusion_match.groups()
            exclusion_info = {
                'rider_code': rider_code,
                'filter_type': filter_type,
                'description': description,
                'result': result,
                'reason': reason
            }
            print(f"找到目标骑手排除信息: {exclusion_info}")

        # 3. 收集成本排名信息
        cost_ranking = None
        if cost_info:
            sorted_riders = sorted(cost_info.items(), key=lambda x: x[1]['cost'])
            cost_ranking = {
                'total_riders': len(sorted_riders),
                'sorted_list': sorted_riders,
                'lowest_cost': sorted_riders[0] if sorted_riders else None,
                'target_rank': None
            }
            
            # 找到目标骑手的排名
            for i, (r_id, _) in enumerate(sorted_riders, 1):
                if r_id == rider_id:
                    cost_ranking['target_rank'] = i
                    break

        # 4. 开始构造完整的分析报告
        result = [f"=== 骑手{rider_id}完整信息分析 ===", ""]
        
        # 4.1 基本信息
        if target_cost_info:
            result.extend([
                "【基本信息】",
                f"- 骑手ID: {rider_id}",
                f"- 骑手代码: {target_cost_info['rider_code']}",
                f"- 总成本: {target_cost_info['cost_str']}",
                f"- 成本排名: 第{target_cost_info['rank']}位",
                ""
            ])
        else:
            result.extend([
                "【基本信息】",
                f"- 骑手ID: {rider_id}",
                "- 成本信息: 未在配送成本日志中找到",
                ""
            ])

        # 4.2 排除原因（如果有）
        if exclusion_info:
            result.extend([
                "【排除原因】",
                f"- 过滤器类型: {exclusion_info['filter_type']}",
                f"- 过滤器描述: {exclusion_info['description']}",
                f"- 检查结果: {exclusion_info['result']}",
                f"- 具体原因: {exclusion_info['reason']}",
                ""
            ])
        else:
            result.extend([
                "【排除原因】",
                f"- 在订单35863807的排除列表中未找到骑手{rider_id}的记录",
                "- 可能原因: 1)未参与此订单派送 2)通过了所有过滤器但成本较高 3)其他业务逻辑排除",
                ""
            ])

        # 4.3 成本对比分析
        if cost_ranking:
            result.extend([
                "【成本对比分析】",
                f"- 参与成本计算的骑手总数: {cost_ranking['total_riders']}人"
            ])
            
            if cost_ranking['target_rank']:
                result.append(f"- 目标骑手成本排名: 第{cost_ranking['target_rank']}/{cost_ranking['total_riders']}位")
            
            if cost_ranking['lowest_cost']:
                lowest_id, lowest_info = cost_ranking['lowest_cost']
                result.extend([
                    f"- 成本最低骑手: {lowest_id}({lowest_info['rider_code']})，成本{lowest_info['cost_str']}",
                ])
                
                if target_cost_info and lowest_id != rider_id:
                    cost_diff = target_cost_info['cost'] - lowest_info['cost']
                    result.append(f"- 成本差距: 比最低成本高出 {cost_diff}")

            result.extend([
                "",
                "成本排名前5位:",
            ])
            for i, (r_id, r_info) in enumerate(cost_ranking['sorted_list'][:5], 1):
                marker = " ← 目标骑手" if r_id == rider_id else ""
                result.append(f"  {i}. 骑手{r_id}({r_info['rider_code']}): {r_info['cost_str']}{marker}")
            result.append("")

        # 4.4 综合结论
        result.extend([
            "【综合结论】"
        ])
        
        if exclusion_info:
            result.append(f"骑手{rider_id}未被选中的主要原因是被{exclusion_info['filter_type']}过滤器排除。")
            result.append(f"具体原因: {exclusion_info['reason']}")
            if target_cost_info:
                result.append(f"即使该骑手的成本为{target_cost_info['cost_str']}（排名第{target_cost_info['rank']}），但由于未通过必要的过滤条件，无法参与最终的成本竞争。")
        elif target_cost_info and cost_ranking and cost_ranking['target_rank'] and cost_ranking['target_rank'] > 1:
            result.append(f"骑手{rider_id}未被选中的主要原因是成本较高。")
            result.append(f"该骑手成本{target_cost_info['cost_str']}，在{cost_ranking['total_riders']}个候选骑手中排名第{cost_ranking['target_rank']}。")
            result.append("根据成本优先原则，系统选择了成本更低的骑手。")
        else:
            result.append(f"基于当前信息，无法确定骑手{rider_id}未被选中的具体原因。")
            result.append("建议检查更详细的派单日志或联系技术支持。")

        return "\n".join(result)

    except Exception as e:
        return f"分析成本时出错: {e}"

def query_rider_info(user_question, file_path="e:\\llm\\03-langchain-rag\\data\\CA7044794619转换后.txt"):
    """
    根据用户问题查询骑手相关信息
    
    Args:
        user_question (str): 用户的问题，例如："为什么21103这个骑手没有被选中"
        file_path (str): 文档路径
    
    Returns:
        str: 分析结果
    """
    try:
        # 判断是否为成本相关问题
        cost_keywords = ["成本", "总成本", "费用", "为什么没有被选中", "为什么没被选中", "没有选中"]
        if any(keyword in user_question for keyword in cost_keywords):
            print("检测到成本相关问题，使用成本分析...")
            return analyze_rider_cost(user_question, file_path)
        
        # 原有的过滤器分析逻辑
        documents, splits = load_document_and_split(file_path)

        # 向量化并在内存中加载
        vectorstore = Chroma.from_documents(
            documents=splits,
            embedding=embedding,
            collection_metadata={"hnsw:space": "cosine"}
        )
        print("文档向量化完成，已在内存中加载")
        
        # 从用户问题中提取骑手ID
        rider_id_match = re.search(r'\d+', user_question)
        if rider_id_match:
            rider_id = rider_id_match.group()
            print(f"检测到骑手ID: {rider_id}")
            
            # 构造多个查询策略（增强：加入 S:{rider_id} 及空格变体）
            queries = [
                f"骑手{rider_id}",
                f"骑手 {rider_id}",
                f"S:{rider_id}",
                f"排除 骑手 {rider_id}",
                f"{rider_id} 过滤",
                user_question,  # 原始用户问题
                f"{rider_id} AvailableRiderDistanceFilter ServiceAreaFilter CashFilter"  # 常见过滤器
            ]
        else:
            # 如果没有找到骑手ID，直接使用用户问题
            queries = [user_question]
        
        print(f"生成的查询策略: {queries}")
        
        all_docs = []
        for query in queries:
            docs = vectorstore.similarity_search(query, k=5)
            all_docs.extend(docs)
            print(f"\n查询: {query}")
            print(f"检索到 {len(docs)} 个相关文档")
        
        # 去重并合并结果
        unique_docs = []
        seen_content = set()
        for doc in all_docs:
            if doc.page_content not in seen_content:
                unique_docs.append(doc)
                seen_content.add(doc.page_content)
        
        print(f"\n总计检索到 {len(unique_docs)} 个唯一文档")
        
        # 优先保留明确包含骑手ID的文档
        if rider_id_match:
            rider_id = rider_id_match.group()
            mention_docs = [
                d for d in unique_docs
                if (rider_id in d.page_content) or (f"S:{rider_id}" in d.page_content) or (f"骑手 {rider_id}" in d.page_content)
            ]
            print(f"明确包含骑手ID的文档数量: {len(mention_docs)}")
            
            # 兜底：如果相似度检索未命中包含 rider_id 的文本，改为关键字扫描原始 splits
            if len(mention_docs) == 0:
                print("未在相似度检索结果中找到明确包含该骑手ID的文本，启用关键字扫描兜底...")
                mention_docs = [
                    s for s in splits
                    if (rider_id in s.page_content) or (f"S:{rider_id}" in s.page_content) or (f"骑手 {rider_id}" in s.page_content) or (f"骑手{rider_id}" in s.page_content)
                ]
                print(f"关键字扫描命中文档数量: {len(mention_docs)}")
        else:
            mention_docs = unique_docs
        
        # 将检索结果交给LLM分析
        context = "\n\n".join([doc.page_content for doc in (mention_docs if len(mention_docs) > 0 else unique_docs)])
        
        # 构造提示词
        prompt = f"""根据以下文档内容，请详细分析用户的问题：

用户问题：{user_question}

文档内容：
{context}

请仔细查找文档中的相关信息，并详细分析回答用户的问题。如果涉及骑手过滤，请说明具体的过滤器类型和原因。"""
        
        print("\n正在让LLM分析...")
        response = llm.invoke(prompt)
        result = response.content
        print(f"\n分析结果：\n{result}")
            
        return result
        
    except Exception as e:
        error_msg = f"处理查询时出错: {e}"
        print(error_msg)
        return error_msg

def main():
    """主函数：演示如何使用查询功能"""
    # 示例查询
    test_questions = [
        "为什么21103这个骑手没有被选中",
        "骑手19523为什么没有接到订单",
        "骑手14885为什么没有被选中",  # 成本相关问题
        "订单35863807的派送情况"
    ]
    
    print("=== RAG 骑手信息查询系统 ===")
    
    # 可以从命令行参数获取问题，或者使用预设的测试问题
    import sys
    if len(sys.argv) > 1:
        user_question = " ".join(sys.argv[1:])
        print(f"\n用户问题: {user_question}")
        result = query_rider_info(user_question)
        if result:
            print(f"\n问题：{user_question}")
            print(f"答案：{result}")
    else:
        # 演示模式：收集所有问题和答案，最后统一输出
        results = []
        for i, question in enumerate(test_questions, 1):
            print(f"正在处理问题 {i}: {question}")
            result = query_rider_info(question)
            results.append((i, question, result if result else "未能获取答案"))
        
        # 最终统一输出所有问题和答案
        print("\n" + "="*80)
        print("查询结果汇总")
        print("="*80)
        for i, question, answer in results:
            print(f"\n问题{i}：{question}")
            print(f"问题{i}答案：{answer}")
            if i < len(results):
                print("-" * 60)

if __name__ == "__main__":
    main()
